#pragma once

#include <geo/geo.h>
#include <mathlib/mathlib.h>
#include <matrix/math.hpp>

#include "common.h"
#include "utils.hpp"

using matrix::AxisAnglef;
using matrix::Dcmf;
using matrix::Eulerf;
using matrix::Matrix3f;
using matrix::Quatf;
using matrix::Vector2f;
using matrix::Vector3f;
using matrix::wrap_pi;

static constexpr uint8_t N_MODELS_EKFGSF = 5;

// Required math constants
static constexpr float _m_2pi_inv = 0.159154943f;
static constexpr float _m_pi = 3.14159265f;
static constexpr float _m_pi2 = 1.57079632f;

using namespace estimator;

class EKFGSF_yaw {
public:
    EKFGSF_yaw();

    // Update Filter States - this should be called whenever new IMU data is available
    void update(const imuSample& imu_sample,
                bool run_EKF,                   // set to true when flying or movement is suitable for yaw estimation
                float airspeed,                 // true airspeed used for centripetal accel compensation - set to 0 when not required.
                const Vector3f& imu_gyro_bias); // estimated rate gyro bias (rad/sec)

    void setVelocity(const Vector2f& velocity, // NE velocity measurement (m/s)
                     float accuracy);          // 1-sigma accuracy of velocity measurement (m/s)

    // get solution data for logging
    bool getLogData(float* yaw_composite,
                    float* yaw_composite_variance,
                    float yaw[N_MODELS_EKFGSF],
                    float innov_VN[N_MODELS_EKFGSF],
                    float innov_VE[N_MODELS_EKFGSF],
                    float weight[N_MODELS_EKFGSF]) const;

    // get yaw estimate and the corresponding variance
    // return false if no yaw estimate available
    bool getYawData(float* yaw, float* yaw_variance) const;

private:
    // Parameters - these could be made tuneable
    const float _gyro_noise { 1.0e-1f };   // yaw rate noise used for covariance prediction (rad/sec)
    const float _accel_noise { 2.0f };     // horizontal accel noise used for covariance prediction (m/sec**2)
    const float _tilt_gain { 0.2f };       // gain from tilt error to gyro correction for complementary filter (1/sec)
    const float _gyro_bias_gain { 0.04f }; // gain applied to integral of gyro correction for complementary filter (1/sec)

    // Declarations used by the bank of N_MODELS_EKFGSF AHRS complementary filters

    Vector3f _delta_ang {};  // IMU delta angle (rad)
    Vector3f _delta_vel {};  // IMU delta velocity (m/s)
    float _delta_ang_dt {};  // _delta_ang integration time interval (sec)
    float _delta_vel_dt {};  // _delta_vel integration time interval (sec)
    float _true_airspeed {}; // true airspeed used for centripetal accel compensation (m/s)

    struct _ahrs_ekf_gsf_struct {
        Dcmf R;             // matrix that rotates a vector from body to earth frame
        Vector3f gyro_bias; // gyro bias learned and used by the quaternion calculation
        bool aligned;       // true when AHRS has been aligned
        float vel_NE[2];    // NE velocity vector from last GPS measurement (m/s)
        bool fuse_gps;      // true when GPS should be fused on that frame
        float accel_dt;     // time step used when generating _simple_accel_FR data (sec)
    } _ahrs_ekf_gsf[N_MODELS_EKFGSF] {};

    bool _ahrs_ekf_gsf_tilt_aligned {}; // true the initial tilt alignment has been calculated
    float _ahrs_accel_fusion_gain {};   // gain from accel vector tilt error to rate gyro correction used by AHRS calculation
    Vector3f _ahrs_accel {};            // low pass filtered body frame specific force vector used by AHRS calculation (m/s/s)
    float _ahrs_accel_norm {};          // length of _ahrs_accel specific force vector used by AHRS calculation (m/s/s)

    // calculate the gain from gravity vector misalingment to tilt correction to be used by all AHRS filters
    float ahrsCalcAccelGain() const;

    // update specified AHRS rotation matrix using IMU and optionally true airspeed data
    void ahrsPredict(const uint8_t model_index);

    // align all AHRS roll and pitch orientations using IMU delta velocity vector
    void ahrsAlignTilt();

    // align all AHRS yaw orientations to initial values
    void ahrsAlignYaw();

    // Efficient propagation of a delta angle in body frame applied to the body to earth frame rotation matrix
    Matrix3f ahrsPredictRotMat(const Matrix3f& R, const Vector3f& g);

    // Declarations used by a bank of N_MODELS_EKFGSF EKFs

    struct _ekf_gsf_struct {
        matrix::Vector3f X;                       // Vel North (m/s),  Vel East (m/s), yaw (rad)s
        matrix::SquareMatrix<float, 3> P;         // covariance matrix
        matrix::SquareMatrix<float, 2> S_inverse; // inverse of the innovation covariance matrix
        float S_det_inverse;                      // inverse of the innovation covariance matrix determinant
        matrix::Vector2f innov;                   // Velocity N,E innovation (m/s)
    } _ekf_gsf[N_MODELS_EKFGSF] {};

    bool _vel_data_updated {};         // true when velocity data has been updated
    bool _run_ekf_gsf {};              // true when operating condition is suitable for to run the GSF and EKF models and fuse velocity data
    Vector2f _vel_NE {};               // NE velocity observations (m/s)
    float _vel_accuracy {};            // 1-sigma accuracy of velocity observations (m/s)
    bool _ekf_gsf_vel_fuse_started {}; // true when the EKF's have started fusing velocity data and the prediction and update processing is active

    // initialise states and covariance data for the GSF and EKF filters
    void initialiseEKFGSF();

    // predict state and covariance for the specified EKF using inertial data
    void predictEKF(const uint8_t model_index);

    // update state and covariance for the specified EKF using a NE velocity measurement
    // return false if update failed
    bool updateEKF(const uint8_t model_index);

    inline float sq(float x) const { return x * x; };

    // Declarations used by the Gaussian Sum Filter (GSF) that combines the individual EKF yaw estimates

    matrix::Vector<float, N_MODELS_EKFGSF> _model_weights {};
    float _gsf_yaw {};          // yaw estimate (rad)
    float _gsf_yaw_variance {}; // variance of yaw estimate (rad^2)

    // return the probability of the state estimate for the specified EKF assuming a gaussian error distribution
    float gaussianDensity(const uint8_t model_index) const;
};
